{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import os\n",
    "import glog as log\n",
    "import copy\n",
    "\n",
    "from __future__ import division\n",
    "\n",
    "import open3d as o3d\n",
    "from open3d import JVisualizer\n",
    "import pandas as pd\n",
    "\n",
    "from evaluation.tools.mesh import Mesh\n",
    "from evaluation.tools.mesh_evaluator import MeshEvaluator\n",
    "\n",
    "# Rotation matrices:\n",
    "# East North Up (ENU) frame to Unity's world frame of reference\n",
    "enu_R_unity = np.array([[1, 0, 0],\n",
    "                        [0, 0, 1],\n",
    "                        [0, 1, 0]])\n",
    "unity_R_enu = np.transpose(enu_R_unity)\n",
    "\n",
    "# Right Handed frame to Unity's Left Handed frame of reference\n",
    "righthand_R_lefthand = np.array([[1, 0, 0],\n",
    "                                 [0, -1, 0],\n",
    "                                 [0, 0, 1]])\n",
    "lefthand_R_righthand = np.transpose(righthand_R_lefthand)\n",
    "\n",
    "\n",
    "visualize = False"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# FILL PATHS BELOW\n",
    "#gt_mesh_path = \"/home/tonirv/Downloads/tesse_multiscene_office1_3d_semantic_v5.ply\"\n",
    "#est_mesh_path = \"/home/tonirv/Downloads/tesse_semantics_2.ply\"\n",
    "\n",
    "#gt_mesh_path = \"/home/tonirv/Code/ROS/flight_goggles_ws/src/voxblox/voxblox_ros/mesh_results/semantic_mesh_tonirv_ld_9118_6487309760727328010.ply\"\n",
    "#est_mesh_path = \"/home/tonirv/Code/ROS/flight_goggles_ws/src/voxblox/voxblox_ros/mesh_results/semantic_mesh_tonirv_ld_9118_6487309760727328010.ply\"\n",
    "\n",
    "#gt_mesh_path = \"/home/tonirv/Downloads/tesse_multiscene_office1_3d_semantic_v5.ply\"\n",
    "#est_mesh_path = \"/home/tonirv/Code/ROS/flight_goggles_ws/src/voxblox/voxblox_ros/mesh_results/tesse_semantics_3.ply\"\n",
    "\n",
    "gt_mesh_path = \"/home/tonirv/Downloads/office1_tony.ply\"\n",
    "\n",
    "est_base_path = \"/home/tonirv/Code/ROS/kimera_ws/src/Kimera-Semantics/kimera_semantics_ros/mesh_results/\"\n",
    "#est_mesh_names = [\"humans_6_long_gt_dyn.ply\",\n",
    "#                  \"humans_6_long_gt_no_dyn.ply\",\n",
    "#                  \"humans_12_long_gt_dyn.ply\",\n",
    "#                  \"humans_12_long_gt_no_dyn.ply\",\n",
    "#                  \"humans_30_long_gt_dyn.ply\",\n",
    "#                  \"humans_30_long_gt_no_dyn.ply\"\n",
    "#                 ]\n",
    "est_mesh_names = [\"humans_6_longvio_dyn.ply\",\n",
    "  \"humans_6_longvio_no_dyn.ply\",\n",
    "  #\"humans_12_long_vio_dyn.ply\",\n",
    "  #\"humans_12_long_vio_no_dyn.ply\",\n",
    "  #\"humans_30_long_vio_dyn.ply\",\n",
    "  #\"humans_30_long_vio_no_dyn.ply\"\n",
    " ]\n",
    "\n",
    "est_mesh_paths = []\n",
    "for est_mesh_name in est_mesh_names:\n",
    "    est_mesh_path = est_base_path + est_mesh_name\n",
    "    est_mesh_paths.append(est_mesh_path)\n",
    "    print(est_mesh_path)\n",
    "\n",
    "\n",
    "#est_mesh_path = \"/home/tonirv/Code/ROS/kimera_ws/src/Kimera-Semantics/kimera_semantics_ros/mesh_results/humans_6_long_gt_no_dyn.ply\"\n",
    "\n",
    "#est_mesh_path = \"/home/tonirv/Code/ROS/kimera_ws/src/Kimera-Semantics/kimera_semantics_ros/mesh_results/humans_12_long_gt_dyn.ply\"\n",
    "#est_mesh_path = \"/home/tonirv/Code/ROS/kimera_ws/src/Kimera-Semantics/kimera_semantics_ros/mesh_results/humans_12_long_gt_no_dyn.ply\"\n",
    "\n",
    "#est_mesh_path = \"/home/tonirv/Code/ROS/kimera_ws/src/Kimera-Semantics/kimera_semantics_ros/mesh_results/humans_12_long_gt_dyn.ply\"\n",
    "#est_mesh_path = \"/home/tonirv/Code/ROS/kimera_ws/src/Kimera-Semantics/kimera_semantics_ros/mesh_results/humans_12_long_gt_no_dyn.ply\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(\"Loading Ground-truth mesh...\")\n",
    "gt_mesh_original = Mesh(gt_mesh_path)\n",
    "\n",
    "# Transform Meshes to same frame of reference\n",
    "gt_mesh = copy.deepcopy(gt_mesh_original)\n",
    "\n",
    "# Align Pointclouds Manually:\n",
    "#est_mesh.mesh_o3d.translate([0, -5, 0])\n",
    "#gt_mesh.transform_left(righthand_R_lefthand)\n",
    "gt_mesh.transform_left(enu_R_unity)\n",
    "\n",
    "if(visualize):\n",
    "    vis = o3d.visualization.Visualizer()\n",
    "    vis.create_window()\n",
    "    vis.get_render_option().mesh_show_back_face = True\n",
    "    vis.add_geometry(est_mesh.mesh_o3d)\n",
    "    vis.add_geometry(gt_mesh.mesh_o3d)\n",
    "    vis.add_geometry(o3d.geometry.create_mesh_coordinate_frame(size=4))\n",
    "    vis.run()\n",
    "    vis.destroy_window()\n",
    "\n",
    "NUMBER_OF_SAMPLES=1000000\n",
    "gt_pcl = o3d.geometry.sample_points_uniformly(gt_mesh.mesh_o3d, NUMBER_OF_SAMPLES)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(\"Loading Estimated meshes...\")\n",
    "dict_est_pcls = dict()\n",
    "for est_mesh_path in est_mesh_paths:\n",
    "    est_mesh_original = Mesh(est_mesh_path)\n",
    "    est_mesh = copy.deepcopy(est_mesh_original)\n",
    "    # Don't sample estimated mesh, just pick vertices, otw you'll be mixing colors...\n",
    "    # est_pcl = o3d.geometry.sample_points_uniformly(est_mesh.mesh_o3d, NUMBER_OF_SAMPLES)\n",
    "    est_pcl = o3d.io.read_point_cloud(est_mesh_path)\n",
    "    dict_est_pcls[est_mesh_path] = (est_pcl)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# ICP\n",
    "def draw_registration_result(source, target, transformation):\n",
    "    source_temp = copy.deepcopy(source)\n",
    "    target_temp = copy.deepcopy(target)\n",
    "    source_temp.paint_uniform_color([1, 0.706, 0])\n",
    "    target_temp.paint_uniform_color([0, 0.651, 0.929])\n",
    "    source_temp.transform(transformation)\n",
    "    o3d.visualization.draw_geometries([source_temp, target_temp])\n",
    "def draw_correspondences(source, target, correspondences):\n",
    "    source_temp = copy.deepcopy(source)\n",
    "    target_temp = copy.deepcopy(target)\n",
    "    #source_temp.paint_uniform_color([1, 0.706, 0])\n",
    "    #target_temp.paint_uniform_color([0, 0.651, 0.929])\n",
    "    o3d.visualization.draw_geometries([source_temp, #target_temp, \n",
    "                                       correspondences])\n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# ICP params\n",
    "ICP_THRESHOLD = 1.5\n",
    "trans_init = np.asarray([[1.0, 0.0, 0.0, 0.0],\n",
    "                         [0.0, 1.0, 0.0, 0.0],\n",
    "                         [0.0, 0.0, 1.0, 0.0],\n",
    "                         [0.0, 0.0, 0.0, 1.0]])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def visualize_two_pcls(gt_pcl, est_pcl):\n",
    "    vis = o3d.visualization.Visualizer()\n",
    "    vis.create_window()\n",
    "    vis.get_render_option().mesh_show_back_face = True\n",
    "    vis.add_geometry(gt_pcl)\n",
    "    vis.add_geometry(est_pcl)\n",
    "    vis.run()\n",
    "    vis.destroy_window()  "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from tqdm import tqdm, trange\n",
    "from random import random, randint\n",
    "from time import sleep\n",
    "\n",
    "def evaluate_icp_for_clouds(gt_pcl, dict_est_pcls):\n",
    "    with trange(len(dict_est_pcls)) as t:\n",
    "        i = 0\n",
    "        for est_pcl_key in dict_est_pcls:\n",
    "            # Description will be displayed on the left\n",
    "            t.set_description('GEN %i' % i)\n",
    "            # Postfix will be displayed on the right,\n",
    "            # formatted automatically based on argument's datatype\n",
    "            t.set_postfix(loss=random(), gen=randint(1,999), str='h',\n",
    "                          lst=[1, 2])\n",
    "\n",
    "            est_pcl = dict_est_pcls[est_pcl_key]\n",
    "            if (visualize):\n",
    "                visualize_two_pcls(gt_pcl, est_pcl)\n",
    "                draw_registration_result(est_pcl, gt_pcl, trans_init)\n",
    "\n",
    "            evaluation = o3d.registration.evaluate_registration(est_pcl, gt_pcl, ICP_THRESHOLD, trans_init)\n",
    "            reg_p2p = o3d.registration.registration_icp(\n",
    "            est_pcl, gt_pcl, ICP_THRESHOLD, trans_init,\n",
    "            o3d.registration.TransformationEstimationPointToPoint(),\n",
    "            o3d.registration.ICPConvergenceCriteria(max_iteration = 2))\n",
    "            correspondences = reg_p2p.correspondence_set\n",
    "\n",
    "            if (visualize):\n",
    "                # Draw Registration Result\n",
    "                draw_registration_result(est_pcl, gt_pcl, reg_p2p.transformation)\n",
    "\n",
    "            print(\"# # # # REGISTRATION INLIER RMSE for: %s \" % est_pcl_key)\n",
    "            print(reg_p2p.inlier_rmse)\n",
    "            print(\"\")\n",
    "            i = i + 1\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# RUN FULL EVALUATION\n",
    "evaluate_icp_for_clouds(gt_pcl, dict_est_pcls)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Visualize initial registration problem\n",
    "if (visualize):\n",
    "    draw_registration_result(est_pcl, gt_pcl, trans_init)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Evaluate current fit between pointclouds\n",
    "evaluation = o3d.registration.evaluate_registration(est_pcl, gt_pcl, ICP_THRESHOLD, trans_init)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(\"Initial registration\")\n",
    "print(evaluation)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(\"Apply point-to-point ICP\")\n",
    "reg_p2p = o3d.registration.registration_icp(\n",
    "    est_pcl, gt_pcl, ICP_THRESHOLD, trans_init,\n",
    "    o3d.registration.TransformationEstimationPointToPoint(),\n",
    "    o3d.registration.ICPConvergenceCriteria(max_iteration = 2000))\n",
    "correspondences = reg_p2p.correspondence_set"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(reg_p2p)\n",
    "print(\"\")\n",
    "\n",
    "print(\"Transformation is:\")\n",
    "print(reg_p2p.transformation)\n",
    "print(\"\")\n",
    "\n",
    "print(\"Correspondence Set:\")\n",
    "print(reg_p2p.correspondence_set)\n",
    "print(\"\")\n",
    "\n",
    "print(\"# # # # REGISTRATION INLIER RMSE: # # # # \")\n",
    "print(reg_p2p.inlier_rmse)\n",
    "print(\"\")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
